# -*- coding:utf-8 -*-
"""
@author: fsksf

@since: 2022/2/8 19:00
"""
import os
from tests import TESTS_DIR
from typing import List
from enum import Enum
from datetime import datetime
import pandas as pd
from vcat.interface import BaseHistoryFeeder
from vcat.core.constant import DataLevel, Exchange, ContractType

LcCsvExchange = Enum('LcCsvExchange',
                     {
                         Exchange.SHEX.name: 'XSHG',
                         Exchange.SZEX.name: 'XSHE',
                         Exchange.DCE.name: 'XDCE'
                     })

LcCsvContractType = Enum('LcCsvContractType',
                         {
                             ContractType.STK.name: 'stock',
                             ContractType.INX.name: 'index',
                             ContractType.FUT.name: 'futures'
                         })


def to_csv_exchange(exchange: Exchange):
    return LcCsvExchange[exchange.name]


def from_csv_exchange(exchange_str: str):
    return Exchange[LcCsvExchange(exchange_str)]


def to_csv_contract_type(con_type: ContractType):
    return LcCsvContractType[con_type.name]


def from_csv_contract_type(con_type_str: str):
    return ContractType[LcCsvContractType(con_type_str).name]


class CSVHistoryFeeder(BaseHistoryFeeder):

    def __init__(self):
        self._minute_bar = pd.read_csv(os.path.join(TESTS_DIR, 'data/quote/bar_minute.csv'),
                                       parse_dates=['time']
                                       )
        self._daily_bar = pd.read_csv(os.path.join(TESTS_DIR, 'data/quote/bar_daily.csv'),
                                      parse_dates=['time']
                                      )

    def login(self, params):
        return True

    def get_history(self, contract_ids: List[str], start_dt, end_dt, count=None, freq: DataLevel=DataLevel.minute):
        if start_dt is None and count is None:
            return 
        if end_dt is None:
            end_dt = datetime.now(tz=None)
        df_list = []
        if freq == DataLevel.minute:
            for cid in contract_ids:
                code, exchange, c_typ = cid.rsplit('.', 2)
                csv_exchange = to_csv_exchange(Exchange[exchange]).value
                symbol = f'{code}.{csv_exchange}'
                mask = (self._minute_bar['code'] == symbol) & (self._minute_bar['time'] <= end_dt)
                if count:
                    df = self._minute_bar[mask].iloc[-count:]
                else:
                    df = self._minute_bar[
                        mask & (self._minute_bar['time'] >= start_dt)
                    ]
                df['contract_type'] = c_typ
                df['code'] = code
                df['exchange'] = exchange
                df_list.append(df)
        elif freq == DataLevel.day:
            for cid in contract_ids:
                code, exchange, c_typ = cid.rsplit('.', 2)
                csv_exchange = to_csv_exchange(Exchange[exchange]).value
                symbol = f'{code}.{csv_exchange}'
                mask = (self._daily_bar['code'] == symbol) & (self._daily_bar['time'] <= end_dt)
                if count:
                    df = self._daily_bar[
                        mask
                    ].iloc[-count:]
                else:
                    df = self._daily_bar[
                        mask & (self._daily_bar['time'] >= start_dt)
                    ]
                df['contract_type'] = c_typ
                df['code'] = code
                df['exchange'] = exchange
                df_list.append(df)
        else:
            return None
        out = pd.concat(df_list, axis=0)
        out.rename(columns={'open': 'open_price', 'high': 'high_price', 'low': 'low_price', 'close': 'close_price',
                            'time': 'dt'}, inplace=True)
        return out
